#!/usr/bin/env python
# coding: utf-8

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import warnings
from matplotlib.ticker import MultipleLocator
from matplotlib.patches import Rectangle  # Import Rectangle
class WaferMap:

    def __init__(self):
        self.sample_size = None
        self.die_pitch = None
        self.die_origin = None
        self.center_location = None
        self.defect_list = pd.DataFrame()

    def load_klarf(self, file, name=None):

        with open(file) as f:
            data = [i.strip() for i in f.readlines()]
        cols = []
        defects = []
        defect_collect = False
        for i in data:
            
            if "SampleSize" in i:
                if self.sample_size and self.sample_size != (float(i[:-1].split()[-1]) * 1000):
                    warnings.warn(f"Warning: Files contain inconsistant sample sizes! {self.sample_size/1000}mm vs {(float(i[:-1].split()[-1]))}mm!")
                # i的值为SampleSize 1 150;, float(i[:-1].split()[-1])的值为150   *1000的原因未知
                # 150000
                self.sample_size = float(i[:-1].split()[-1]) * 1000
            
            elif "DiePitch" in i:
                if self.die_pitch and self.die_pitch != [float(j) for j in i[:-1].split(" ")[1:]]:
                    warnings.warn("Warning: Files contain inconsistent die sizes. Recommend setting die_line_alpha to 0 for plotting")
                self.die_pitch = [float(j) for j in i[:-1].split(" ")[1:]]  #self.die_pitch的值为 1695.1 2493.0
            
            elif "DieOrigin" in i:
                self.die_origin = [float(j) for j in i[:-1].split(" ")[1:]] # self.die_origin的值为 0 0
            
            elif "CenterLocation" in i:
                self.center_location = [float(j) for j in i[:-1].split(" ")[1:]] # 445.157000000007 739.754000000001
            
            elif "DefectRecordSpec" in i:
                for j in i[:-1].split()[2::]:
                    cols.append(j)  # DefectRecordSpec 11 DEFECTID X Y XREL YREL XINDEX YINDEX CLASSNUMBER TEST IMAGECOUNT IMAGELIST
            
# 00:
# 'DEFECTID'
# 01:
# 'X'
# 02:
# 'Y'
# 03:
# 'XREL'
# 04:
# 'YREL'
# 05:
# 'XINDEX'
# 06:
# 'YINDEX'
# 07:
# 'CLASSNUMBER'
# 08:
# 'TEST'
# 09:
# 'IMAGECOUNT'
# 10:
# 'IMAGELIST'
            elif "DefectList" in i:
                defect_collect = True

            elif "SummarySpec" in i:
                defect_collect = False

            elif defect_collect:
                if "TiffFileName" in i or i =="":
                    continue
                else:
                    first_11_values = i.split()[:13]
                    defects.append(first_11_values)
# DefectRecordSpec 11 DEFECTID X                  Y               XREL            YREL                XINDEX YINDEX CLASSNUMBER TEST IMAGECOUNT IMAGELIST
# 1                            39466.1765775107 41130.0826670413 478.876577510746 1242.08266704126    23      17      113         1   1           1 1 0;                    
# 00:
# 'DefectRecordSpec'
# 01:
# '11'
# 02:
# 'DEFECTID'
# 03:
# 'X'
# 04:
# 'Y'
# 05:
# 'XREL'
# 06:
# 'YREL'
# 07:
# 'XINDEX'
# 08:
# 'YINDEX'
# 09:
# 'CLASSNUMBER'
# 10:
# 'TEST'
# 11:
# 'IMAGECOUNT'
# 12:
# 'IMAGELIST'
# 13:
# '_XACTUAL'
# 14:
# '_YACTUAL'
# 15:
# '_KLARFNAME'
        df = pd.DataFrame(defects, columns=cols)
        # df = pd.DataFrame(defects, columns=cols).astype(float)
        for col in df.columns:
            df[col] = df[col].apply(convert_to_float)
        print(df['XINDEX'])
        print(df['XREL'])
        print(self.die_pitch)
        print(self.center_location)
        df['_XACTUAL'] = (df['XINDEX'] * self.die_pitch[0]) + df['XREL'] - self.center_location[0]
        df['_YACTUAL'] = (df['YINDEX'] * self.die_pitch[1]) + df['YREL'] - self.center_location[1]

        if name:
            df['_KLARFNAME'] = name
        else:
            df['_KLARFNAME'] = file

        self.defect_list = pd.concat([self.defect_list, df]).reset_index(drop=True)
        print(self.defect_list)

    def plot_wafer_map(self, color='_KLARFNAME', die_line_alpha=0.2, die_line_color='gray', *args, **kwargs):
        # fig1 = plt.figure(figsize=(50.0, 50.0), dpi=100)
        # ax1 = fig1.add_subplot(111, aspect='equal')
        # ax1.xaxis.set_major_locator(MultipleLocator(2000.0))
        # ax1.yaxis.set_major_locator(MultipleLocator(2000.0))
        # width = 4687
        # height = 3645
        # # 画出ng区域
        # ret = plt.Rectangle((0, 0), width, height,color="gray", linestyle='dotted', fill=True, linewidth=1)
        # ax1.add_patch(ret)
        fig, ax = plt.subplots()
        wafer_sin = (np.sin(np.arange(0, 2 * np.pi, 1 / 1000))) * self.sample_size / 2
        wafer_cos = (np.cos(np.arange(0, 2 * np.pi, 1 / 1000))) * self.sample_size / 2
        # Plot the wafer circle
        plt.plot(wafer_sin, wafer_cos, color='k', linewidth=0.3)



            
        for i in range(-int(self.sample_size // self.die_pitch[0]) - 10,
                       int(self.sample_size // self.die_pitch[0]) + 10):

            die_x_lines = -self.center_location[0] + (i * self.die_pitch[0])
            height_squared = (self.sample_size / 2) ** 2 - die_x_lines ** 2
            if height_squared > 0:
                height = np.sqrt(height_squared)

                plt.vlines(die_x_lines,
                           -height,
                           height,
                           color=die_line_color,
                           alpha=die_line_alpha)

        for i in range(-int(self.sample_size // self.die_pitch[1]) - 10,
                       int(self.sample_size // self.die_pitch[1]) + 10):

            die_y_lines = -self.center_location[1] + (i * self.die_pitch[1])
            width_squared = (self.sample_size / 2) ** 2 - die_y_lines ** 2
            if width_squared > 0:
                width = np.sqrt(width_squared)

                plt.hlines(die_y_lines,
                           -width,
                           width,
                           color=die_line_color,
                           alpha=die_line_alpha)
        # print(self.defect_list["XINDEX"])
        # print(self.defect_list["YINDEX"])
        # 这个是红色方框
        all_data = []
        for i in range(len(self.defect_list["XINDEX"])):
            x = self.defect_list["XINDEX"][i]
            y = self.defect_list["YINDEX"][i]
            if (x,y) not in all_data:
                all_data.append((x,y))
        for j in all_data:
            width = self.die_pitch[0]
            height = self.die_pitch[1]
            red_box = Rectangle((j[0]*width, j[1]*height),  width, height, linewidth=1, edgecolor='r', facecolor='none')
            plt.gca().add_patch(red_box)            
        # for i in len
        if color in self.defect_list.columns:
            sns.scatterplot(x=self.defect_list['_XACTUAL'],
                            y=self.defect_list['_YACTUAL'],
                            hue=self.defect_list[color],
                            *args, **kwargs)

            plt.legend(loc='lower left')
        else:
            sns.scatterplot(x=self.defect_list['_XACTUAL'],
                            y=self.defect_list['_YACTUAL'],
                            color=color,
                            *args, **kwargs)


        
        plt.xlim(-self.sample_size / 2 - 2000, self.sample_size / 2 + 2000)
        plt.ylim(-self.sample_size / 2 - 2000, self.sample_size / 2 + 2000)
            # Display the plot
        plt.show()
def convert_to_float(value):
    try:
        return float(value)
    except ValueError:
        return 0  # 或者其他你认为合适的默认值
if __name__=="__main__":
    test_klaf = WaferMap()
    # test_klaf.load_klarf("klarfkit\TEST_12.klarf")
    # test_klaf.load_klarf("D:/QtProject/json_to_klarf/1.klarf")
    klarf_path = r"1.klarf"
    klarf_path = r"D:/QtProject/json_to_klarf/1.klarf"
    test_klaf.load_klarf(klarf_path)
    # test_klaf.load_klarf("D:\workspace\SWDev\SW\Training\Wafer\FileFormats\KLARF\TEST_12.klarf")
    test_klaf.plot_wafer_map()
    # 1（缺陷id） 39466.1765775107（绝对x坐标） 41130.0826670413（绝对Y坐标） 478.876577510746（缺陷在die上的x坐标） 1242.08266704126（缺陷在die上的Y坐标） 23（die的x坐标） 17(die的y坐标) 113 1 1 1 1 0;
    # 